#include <ros/ros.h>
#include <tf/transform_datatypes.h>

#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/Pose.h>
#include <geometry_msgs/TwistStamped.h>
#include <geometry_msgs/Point32.h>
#include <std_msgs/Float64MultiArray.h>
#include <sensor_msgs/PointCloud2.h>
#include <visualization_msgs/Marker.h>
#include <pcl/common/transforms.h>
#include <pcl/common/angles.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/point_types.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/segmentation/extract_clusters.h>
#include <pcl/filters/passthrough.h>
#include <pcl/filters/radius_outlier_removal.h>
#include <iostream>
#include <string>
#include <cmath>
#include <stdlib.h>
#include <mavros_msgs/State.h> 
#include <nav_msgs/Odometry.h>
#include <std_msgs/Float64.h>
#include <queue>
#include <cstdint>
#include <trajectory_msgs/JointTrajectory.h>
#include <std_msgs/ColorRGBA.h>
#include <visualization_msgs/MarkerArray.h>
#include <nav_msgs/Path.h>
#include <trajectory_msgs/JointTrajectoryPoint.h>
#include <algorithm>
#include <gmm_voxel_map/gmm.h>
#include "gmap/GMMCommunication.h"

using namespace std;
using namespace Eigen;


int AGENT_ID= 0;
int COLOR_ID=0;
double init_x=0;
double init_y=0;
double init_z=0;
int count_time = 0; //TEST
double total_time;//TEST

double motor_yaw=0; //TODO
double N_SIGMA = 1.6;

ros::Publisher cloud2_pub;
ros::Publisher cluster_centre_pub;
ros::Publisher ellipsoid_pub;
ros::Publisher ellipsoid_marker_resample_pub;
ros::Publisher ellipsoid_marker_part_pub;
ros::Publisher gmm_pub;
ros::Publisher gmm_mix_pub;
ros::Publisher global_position_pub;
// TEST
ros::Publisher pass_cloud_ros_pub;
ros::Publisher Outlier_cloud_ros_pub;
ros::Publisher resample_pc_pub;
ros::Publisher voxel_pc_pub;
ros::Time time_get_yaw;
ros::Time time_get_pos;

ros::Timer resample_timer;

gmap::GMMCommunication all_gmm_msg[4];
gvm::GMM gmm;

double PI=3.1415926;
bool other_gmm_flag = false;


/**
 * @brief 将点云沿着z轴进行阶段
 * @param limL 最小高度
 * @param limH 最大高度
 */
void PclPassz(pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in, pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_out, double limL, double limH)
{
	pcl::PassThrough<pcl::PointXYZ> pass;
	pass.setInputCloud (cloud_in);
    pass.setFilterFieldName ("z");// 滤波字段设置为z轴方向
    pass.setFilterLimits (limL, limH); 
	// // 设置保留范围内还是过滤掉范围内, 默认为false,可以注释掉。true为过滤掉范围内的，false为过滤掉范围外的
    // pass.setFilterLimitsNegative (false); 
    pass.filter (*cloud_out);
}


/**
 * @brief 体素滤波，下采样
 * @param resolution 分辨率
 */
void PclVoxel(pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in, pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_out,double resolution)
{
	pcl::VoxelGrid<pcl::PointXYZ> sorr;
	sorr.setInputCloud(cloud_in);
	sorr.setLeafSize(resolution, resolution, resolution);
	sorr.filter(*cloud_out);
}


/**
 * @brief 移除异常噪声点云
 * @param cloud_in 
 * @param cloud_out 
 * @param RadiusSearch 设置查询点的半径邻域范围
 * @param MinNeighbors 判断是否为离群点的阈值，即半径内至少包括的点数
 */
void PclOutlier(pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in, pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_out,double RadiusSearch,int MinNeighbors)
{
	pcl::RadiusOutlierRemoval<pcl::PointXYZ> ror;	//创建滤波器对象
	ror.setInputCloud(cloud_in);					//设置待滤波点云
	ror.setRadiusSearch(RadiusSearch); 
	ror.setMinNeighborsInRadius(MinNeighbors);
	//ror.setNegative(true);						//默认false，保存内点；true，保存滤掉的外点
	ror.filter(*cloud_out);
}	


/**
 * @brief 对点云进行欧式聚类
 * 
 * @param cloud_in 
 * @param ClusterTolerance 近邻搜索的半径
 * @param MinClusterSize 最小聚类点数
 */
bool PclEuclideanCluster(pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in,double ClusterTolerance,int MinClusterSize,
						std::vector<Eigen::Vector3d>& means,  std::vector<Eigen::Matrix3d>& var, std::vector<double>& pi_vector,double n_sigma)
{
	pcl::search::KdTree<pcl::PointXYZ>::Ptr tree(new pcl::search::KdTree<pcl::PointXYZ>);
	tree->setInputCloud(cloud_in);
	std::vector<pcl::PointIndices> cluster_indices;
	pcl::EuclideanClusterExtraction<pcl::PointXYZ> ec;
	ec.setClusterTolerance(ClusterTolerance);
	ec.setMinClusterSize(MinClusterSize);
	ec.setMaxClusterSize(999999);	//设置最大聚类点数
	ec.setSearchMethod(tree);
	ec.setInputCloud(cloud_in);
	ec.extract(cluster_indices);

	/// 执行欧式聚类分割，并保存分割结果
	if(cluster_indices.size()==0){
		ROS_INFO("\033[1;%dm In PclEuclideanCluster cluster_indices.size()==0 \033[0m",COLOR_ID);
		return false;
	}
	for (std::vector<pcl::PointIndices>::const_iterator it = cluster_indices.begin(); it != cluster_indices.end(); ++it)
	{
		Eigen::Vector3d cluster_centre=Eigen::Vector3d::Zero();
		double num=0;

		//求取聚类后的中心
		for (std::vector<int>::const_iterator pit = it->indices.begin(); pit != it->indices.end(); pit++)
		{
			Eigen::Vector3d	tmp(cloud_in->points[*pit].x,cloud_in->points[*pit].y,cloud_in->points[*pit].z);
			cluster_centre+=tmp;
			num++;
		}
		cluster_centre/=(double)num;
		means.push_back(cluster_centre);
		pi_vector.push_back((double)num);

		Eigen::Matrix3d var_tmp=Eigen::Matrix3d::Zero();

		for (std::vector<int>::const_iterator pit = it->indices.begin(); pit != it->indices.end(); pit++)
		{
			Eigen::Vector3d	tmp(cloud_in->points[*pit].x,cloud_in->points[*pit].y,cloud_in->points[*pit].z);
			var_tmp+=(tmp-cluster_centre)*(tmp-cluster_centre).transpose();
		}		
		var_tmp = n_sigma * var_tmp / (double)num;
		var.push_back(var_tmp);
	}
	return true;
}


/**
 * @param ALL_GMM true ->TOPIC:/gmm_mix_server，混合后的GMM
 * @param ALL_GMM false->TOPIC:/gmm_server，单个GMM
 */
void sendGmmMsg(std::vector<double> pi_vector,std::vector<Eigen::Vector3d> means, std::vector<Eigen::Matrix3d> var,int n_component,bool ALL_GMM)
{
	gmap::GMMCommunication gmm_msg;
	gmm_msg.agent_id=AGENT_ID;
	gmm_msg.gmm_number=n_component;
	
	for(int i=0;i<n_component;i++)
	{
		gmm_msg.pi.push_back(pi_vector[i]);
		gmm_msg.means.push_back(means[i](0));
		gmm_msg.means.push_back(means[i](1));
		gmm_msg.means.push_back(means[i](2));

		for(int j=0;j<3;j++)
		{
			for(int k=0;k<3;k++)
			{
				gmm_msg.vars.push_back(var[i](j,k));
			}
		}
	}
	gmm_msg.header.stamp=ros::Time::now();
	if(ALL_GMM == true){
		gmm_mix_pub.publish(gmm_msg);
	}
	else{
		gmm_pub.publish(gmm_msg);
		all_gmm_msg[AGENT_ID]=gmm_msg;
	}
	
}

void visual_ellipsoid(std::vector<Eigen::Vector3d> means,std::vector<Eigen::Matrix3d> var,int id,bool resample_flag)
{
    visualization_msgs::MarkerArray marker_array;
    int cluster_num=means.size();
    // ROS_INFO("means var size %ld %ld",means.size(),var.size());

    for(int i=0;i<cluster_num;i++)
    {
        for(int j=0;j<3;j++)
        {
            for(int k=0;k<3;k++)
            {
                if(fabs(var[i](j,k))<1e-5)
                {
                    var[i](j,k)=0;
                }
            }
        }

        EigenSolver<Eigen::Matrix3d> es(var[i]);

        Eigen::Matrix3d value_martrix = es.pseudoEigenvalueMatrix();
        Eigen::Matrix3d vector_martrix = es.pseudoEigenvectors();    
        vector_martrix.col(2)=vector_martrix.col(0).cross(vector_martrix.col(1));
        // std::cout<<"var: "<<var[i]<<endl;

        // std::cout<<"value_martrix: \n"<<value_martrix<<endl;
        // std::cout<<"vector_martrix: \n"<<vector_martrix<<endl;
        // std::cout<<"means[i]: "<<means[i]<<endl;

        double var0=pow(value_martrix(0,0),0.5);
        double var1=pow(value_martrix(1,1),0.5);
        double var2=pow(value_martrix(2,2),0.5);
        // ROS_INFO("i: %d var: %f %f %f, means: %f %f %f",i,var0,var1,var2,means[i](0),means[i](1),means[i](2));

        visualization_msgs::Marker marker;

        marker.header.frame_id = "map";
        marker.header.stamp = ros::Time();
        marker.id = 23123+i+id+AGENT_ID;
        marker.type = visualization_msgs::Marker::SPHERE;
        marker.action = visualization_msgs::Marker::ADD;

        marker.lifetime = ros::Duration(1.0);

        marker.pose.position.x = means[i](0);
        marker.pose.position.y = means[i](1);
        marker.pose.position.z = means[i](2);

        Eigen::Quaterniond quaternion(vector_martrix);
        quaternion.normalize();
        marker.pose.orientation.x = quaternion.x();
        marker.pose.orientation.y = quaternion.y();
        marker.pose.orientation.z = quaternion.z();
        marker.pose.orientation.w = quaternion.w();

        marker.scale.x = var0*2;
        marker.scale.y = var1*2;
        marker.scale.z = var2*2;

        marker.color.a = 0.5; // Don't forget to set the alpha!
        marker.color.r = 0.0;
        marker.color.g = 0.0;
        marker.color.b = 0.0;

        marker_array.markers.push_back(marker);

    }
	if(resample_flag){
		ellipsoid_marker_resample_pub.publish(marker_array);
	}
	else{
		ellipsoid_marker_part_pub.publish(marker_array);
	}
    

    // visulization of cluster centre
    visualization_msgs::Marker point_marker;

    point_marker.id = 333;
    point_marker.header.frame_id = "map";
    point_marker.header.stamp = ros::Time::now();
    point_marker.type = visualization_msgs::Marker::SPHERE_LIST;
    point_marker.action = visualization_msgs::Marker::ADD;
    point_marker.ns = "lines_and_points";

    point_marker.scale.x = 0.4;
    point_marker.scale.y = 0.4;
    point_marker.scale.z = 0.4;

    point_marker.lifetime = ros::Duration(0);

    point_marker.color.r=1;
    point_marker.color.g=0;
    point_marker.color.b=0;
    point_marker.color.a=1.0;

    for(unsigned int i=0;i<means.size();i++)
    {
        geometry_msgs::Point point;
        point.x = means[i](0);
        point.y = means[i](1);
        point.z = means[i](2);
        point_marker.points.push_back(point);  
    }
    cluster_centre_pub.publish(point_marker);

}

void cloudCallback( const sensor_msgs::PointCloud2ConstPtr& cloud)
{
	ros::Time cloudcbstart=ros::Time::now();
    // ROS_INFO("\033[1;%dm CLOUD CALL BACK ENTER ================= \033[0m",COLOR_ID);

    if(cloud->data.size()<=10)return;

    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in(new pcl::PointCloud<pcl::PointXYZ>());
    pcl::fromROSMsg(*cloud, *cloud_in);

	//TEST
	// ros::Time time_after_voxel = ros::Time::now();
	// if(AGENT_ID == 2){
	// 	if(count_time >10){
	// 		total_time += (time_after_voxel - cloud->header.stamp).toSec();
	// 	}
	// 	count_time++;
	// 	if(count_time == 50){
	// 		ROS_INFO(" [GMM]: time_after_voxel-cloud->header.stamp: %f ======== ",total_time/40);
	// 	}
	// }

	// Remove Outliers
	pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_second_filter(new pcl::PointCloud<pcl::PointXYZ>());
	PclOutlier(cloud_in, cloud_second_filter, 0.3, 20); //TODO

	// cout << "cloud_in->points.size()" << cloud_in->points.size() << endl;
	// cout << "cloud_filter->points.size()" << cloud_filter->points.size() << endl;
	// cout << "cloud_voxel->points.size()" << cloud_voxel->points.size() << endl;
	// cout << "cloud_second_filter->points.size()" << cloud_second_filter->points.size() << endl;

	if(cloud_second_filter->points.size()<5){
		ROS_INFO("\033[1;%dm cloud_second_filter < 5 \033[0m",COLOR_ID);
		return;
	}


	std::vector<Eigen::Vector3d> means;
    std::vector<Eigen::Matrix3d> var;
    std::vector<double> pi_vector;

	if (!PclEuclideanCluster(cloud_second_filter, 0.3, 20, means, var, pi_vector,N_SIGMA))
	{
		return;
	}

	// if(AGENT_ID != 0 ){
	visual_ellipsoid(means,var,100,false);  //TODO
	// }
	// ROS_INFO_THROTTLE(3,"\033[1;%dm [WAIT] SENDING GMM MSG ======= \033[0m",COLOR_ID);
	sendGmmMsg(pi_vector, means, var, means.size(), false);

	ros::Time cloudcbend=ros::Time::now();
	// TEST
	// if(AGENT_ID == 1){
    ROS_INFO_THROTTLE(3,"\033[1;%dm CLOUD EXIT, cloud_CB use time: %f ======== \033[0m",COLOR_ID,(cloudcbend-cloudcbstart).toSec());
	// }
}


//接受别人的GMM模型
void gmmCallback( const gmap::GMMCommunication& gmm_msg)
{
	if(gmm_msg.agent_id==AGENT_ID){
		return;
	}
	// ROS_INFO("\033[1;%dm  gmmCallback enter!!!!!!!!!!! \033[0m",COLOR_ID);
	all_gmm_msg[gmm_msg.agent_id]=gmm_msg;
	other_gmm_flag = true;
}

void resample_cb(const ros::TimerEvent &e)
{
	// ROS_INFO("\033[1;%dm  resample enter!!!!!!!!!!! \033[0m",COLOR_ID);
	ros::Time starttime=ros::Time::now();

	if(!other_gmm_flag){
		all_gmm_msg[AGENT_ID].header.stamp = ros::Time::now();
		gmm_mix_pub.publish(all_gmm_msg[AGENT_ID]);
		return;
	}
	
	std::vector<Eigen::Vector3d> means;
    std::vector<Eigen::Matrix3d> var;
    std::vector<double> pi_vector;
	std::vector<Eigen::Vector3d> means_final;
	std::vector<Eigen::Matrix3d> var_final;
	std::vector<double> pi_vector_final;

	for (int uav_i = 0; uav_i < 4;uav_i++){
		gmap::GMMCommunication gmm_msg = all_gmm_msg[uav_i];
		int gmm_number = gmm_msg.gmm_number;
		for (int i = 0; i < gmm_number;i++){
			Eigen::Vector3d tmp;
			tmp << gmm_msg.means[i * 3 + 0], gmm_msg.means[i * 3 + 1], gmm_msg.means[i * 3 + 2];
			means.push_back(tmp);

			Eigen::Matrix3d tmp_var;
			tmp_var << gmm_msg.vars[i * 9 + 0], gmm_msg.vars[i * 9 + 1], gmm_msg.vars[i * 9 + 2],
					   gmm_msg.vars[i * 9 + 3], gmm_msg.vars[i * 9 + 4], gmm_msg.vars[i * 9 + 5],
					   gmm_msg.vars[i * 9 + 6], gmm_msg.vars[i * 9 + 7], gmm_msg.vars[i * 9 + 8];
			var.push_back(tmp_var);

			double tmp_pi=gmm_msg.pi[i];
			pi_vector.push_back(tmp_pi);
		}
	}

	gvm::GMM all_gmm;
	all_gmm.init(pi_vector,means,var,means.size());
	all_gmm.setAgentId(AGENT_ID,COLOR_ID,N_SIGMA);
	pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_resample(new pcl::PointCloud<pcl::PointXYZ>());
	all_gmm.resample(cloud_resample); // TODO
	
	//TEST
	sensor_msgs::PointCloud2 cloud_resample_ros;
    pcl::toROSMsg(*cloud_resample, cloud_resample_ros);
    cloud_resample_ros.header.stamp = ros::Time::now();
    cloud_resample_ros.header.frame_id = "map";
   	resample_pc_pub.publish(cloud_resample_ros);

	// ros::Time endtime=ros::Time::now();
	// ROS_INFO("\033[1;%dm  resample_cb use time: %f\033[0m",COLOR_ID,(endtime-starttime).toSec());

	// VoxelGrid
	pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_voxel(new pcl::PointCloud<pcl::PointXYZ>());
	PclVoxel(cloud_resample, cloud_voxel, 0.1);

	sensor_msgs::PointCloud2 cloud_voxel_ros;
    pcl::toROSMsg(*cloud_voxel, cloud_voxel_ros);
    cloud_voxel_ros.header.stamp = ros::Time::now();
    cloud_voxel_ros.header.frame_id = "map";
   	voxel_pc_pub.publish(cloud_voxel_ros);

	// Remove Outliers
	pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_filter(new pcl::PointCloud<pcl::PointXYZ>());
	PclOutlier(cloud_voxel, cloud_filter, 0.3, 11);

	// TEST
	sensor_msgs::PointCloud2 Outlier_cloud_ros;
    pcl::toROSMsg(*cloud_filter, Outlier_cloud_ros);
    Outlier_cloud_ros.header.stamp = ros::Time::now();
    Outlier_cloud_ros.header.frame_id = "map";
    Outlier_cloud_ros_pub.publish(Outlier_cloud_ros);

	// std::cout << "cloud_resample->points.size()" << cloud_resample->points.size() << std::endl;
	// std::cout << "cloud_voxel->points.size()" << cloud_voxel->points.size() << std::endl;
	
	if(cloud_filter->points.size()<5){
		ROS_INFO("\033[1;%dm cloud_voxel < 5 \033[0m",COLOR_ID);
		return;
	}
	

	if(!PclEuclideanCluster(cloud_filter, 0.3, 10, means_final, var_final, pi_vector_final,N_SIGMA)){
		return;
	}

	sendGmmMsg(pi_vector_final, means_final, var_final, means_final.size(), true);

	// if(AGENT_ID == 0){  //TODO
	// 	visual_ellipsoid(means_final,var_final,100);
	// }
	visual_ellipsoid(means_final,var_final,100,true);

	ros::Time endtime=ros::Time::now();
	ROS_INFO("\033[1;%dm  resample_cb use time: %f\033[0m",COLOR_ID,(endtime-starttime).toSec());
}


void positionCallback(const geometry_msgs::PoseStamped& msg)
{
	geometry_msgs::PoseStamped global_pos_msg;
	global_pos_msg=msg;
	global_pos_msg.pose.position.x = msg.pose.position.x+init_x;
	global_pos_msg.pose.position.y = msg.pose.position.y+init_y;
	global_pos_msg.pose.position.z = msg.pose.position.z+init_z;

	global_position_pub.publish(global_pos_msg);
}


int main (int argc, char** argv) 
{   
    ros::init(argc, argv, "gmm_map_node0");
    ros::NodeHandle nh("~");
	srand((unsigned)time(NULL));

	ros::Duration(5.0).sleep();

	nh.getParam("AGENT_ID",AGENT_ID);
	nh.getParam("x",init_x);
	nh.getParam("y",init_y);
	nh.getParam("z",init_z);
	nh.getParam("N_SIGMA",N_SIGMA);


	COLOR_ID=AGENT_ID+32;

	gmm.setAgentId(AGENT_ID,COLOR_ID,N_SIGMA); // TODO

    const int LOOPRATE = 30;
    ros::Rate loop_rate(LOOPRATE);  

	ros::Subscriber cloud_sub= nh.subscribe("/convert_pointcloud", 1, cloudCallback);
	ros::Subscriber other_gmm_sub = nh.subscribe("/gmm_server", 4, gmmCallback);
    ros::Subscriber position_sub =  nh.subscribe("/mavros/local_position/pose", 1, positionCallback);
	
    cluster_centre_pub= nh.advertise<visualization_msgs::Marker>("/cluster_centr", 1, true);
   	ellipsoid_marker_resample_pub = nh.advertise<visualization_msgs::MarkerArray>("/ellipsoid_resample", 1);
	ellipsoid_marker_part_pub = nh.advertise<visualization_msgs::MarkerArray>("/ellipsoid_part", 1);

	gmm_pub = nh.advertise<gmap::GMMCommunication>("/gmm_server", 1, true); //发布GMM模型
	gmm_mix_pub = nh.advertise<gmap::GMMCommunication>("/gmm_mix_server", 1, true); //发布GMM模型
	global_position_pub= nh.advertise<geometry_msgs::PoseStamped>("/global_position", 1, true);
	
	// TEST
	pass_cloud_ros_pub= nh.advertise<sensor_msgs::PointCloud2>("/pass_pc", 1, true);
	Outlier_cloud_ros_pub= nh.advertise<sensor_msgs::PointCloud2>("/Outlier_pc", 1, true);
	resample_pc_pub= nh.advertise<sensor_msgs::PointCloud2>("/resample_pc", 1, true);
	voxel_pc_pub= nh.advertise<sensor_msgs::PointCloud2>("/voxel_pc", 1, true);

	resample_timer = nh.createTimer(ros::Duration(0.05), resample_cb); //TODO

    // ros::AsyncSpinner spinner(3); // Use 3 threads
    // spinner.start();
    // ros::waitForShutdown();

	ros::spin();

	return 0;
}